"""
Serialization and deserialization for CodeMCP call trees and related data.

Handles conversion between in-memory CallTree/CallTreeNode objects and
database storage models, with support for various serialization formats.
"""

import json
import pickle
import gzip
import hashlib
from typing import Dict, Any, List, Optional, Union, Type
from datetime import datetime
from pathlib import Path

from .models import CallTreeModel, CallTreeNodeModel, AnalysisResultModel
from .database import DatabaseManager
from ..core.base_mcp import CallTree, CallTreeNode
from ..core.error_handler import CodeMCPError, log_debug, log_warning, handle_error


class SerializationError(CodeMCPError):
    """Serialization-specific errors."""
    
    def __init__(self, message: str, operation: str = None, **kwargs):
        super().__init__(message, "SERIALIZATION_ERROR", **kwargs)
        if operation:
            self.details["operation"] = operation


class CallTreeSerializer:
    """
    Handles serialization and deserialization of call trees.
    
    Supports conversion between in-memory CallTree objects and database models,
    as well as various export/import formats.
    """
    
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
        
        # Serialization statistics
        self.stats = {
            "trees_serialized": 0,
            "trees_deserialized": 0,
            "nodes_serialized": 0,
            "nodes_deserialized": 0,
            "serialization_errors": 0
        }
    
    async def serialize_call_tree(self, call_tree: CallTree, 
                                analysis_type: str = "unknown") -> CallTreeModel:
        """
        Serialize a CallTree object to database model.
        
        Args:
            call_tree: In-memory CallTree object
            analysis_type: Type of analysis that created this tree
            
        Returns:
            CallTreeModel instance ready for database storage
        """
        try:
            # Create call tree model
            tree_model = CallTreeModel(
                analysis_type=analysis_type,
                language=call_tree.metadata.get("language", "unknown"),
                project_path=call_tree.metadata.get("project_path"),
                total_nodes=len(call_tree.nodes),
                max_depth=call_tree.calculate_depth(),
                metadata=call_tree.metadata
            )
            
            # Save the tree to get an ID
            tree_model = await self.db_manager.create(tree_model)
            
            # Serialize all nodes
            node_id_mapping = {}  # Map old node IDs to new database IDs
            
            for node_id, node in call_tree.nodes.items():
                node_model = await self.serialize_call_tree_node(node, tree_model.id)
                node_id_mapping[node_id] = node_model.id
            
            # Update node relationships with new IDs
            for node_id, node in call_tree.nodes.items():
                node_model_id = node_id_mapping[node_id]
                node_model = await self.db_manager.get_by_id(CallTreeNodeModel, node_model_id)
                
                # Map caller and callee IDs
                caller_ids = []
                for caller in node.caller_nodes:
                    caller_node_id = f"{caller.file_path}:{caller.line_number}:{caller.name}"
                    if caller_node_id in node_id_mapping:
                        caller_ids.append(node_id_mapping[caller_node_id])
                
                callee_ids = []
                for callee in node.callee_nodes:
                    callee_node_id = f"{callee.file_path}:{callee.line_number}:{callee.name}"
                    if callee_node_id in node_id_mapping:
                        callee_ids.append(node_id_mapping[callee_node_id])
                
                # Update the node model
                await self.db_manager.update(node_model, 
                                           caller_node_ids=caller_ids,
                                           callee_node_ids=callee_ids)
            
            # Set root node if available
            if call_tree.root_node:
                root_node_id = f"{call_tree.root_node.file_path}:{call_tree.root_node.line_number}:{call_tree.root_node.name}"
                if root_node_id in node_id_mapping:
                    await self.db_manager.update(tree_model, 
                                               root_node_id=node_id_mapping[root_node_id])
            
            self.stats["trees_serialized"] += 1
            log_debug(f"Serialized call tree with {len(call_tree.nodes)} nodes")
            
            return tree_model
            
        except Exception as e:
            self.stats["serialization_errors"] += 1
            error_response = handle_error(e, {
                "operation": "serialize_call_tree",
                "analysis_type": analysis_type,
                "node_count": len(call_tree.nodes) if call_tree.nodes else 0
            })
            raise SerializationError(f"Failed to serialize call tree: {error_response['message']}")
    
    async def serialize_call_tree_node(self, node: CallTreeNode, 
                                     call_tree_id: str) -> CallTreeNodeModel:
        """
        Serialize a CallTreeNode to database model.
        
        Args:
            node: In-memory CallTreeNode object
            call_tree_id: ID of the parent call tree
            
        Returns:
            CallTreeNodeModel instance
        """
        try:
            node_model = CallTreeNodeModel(
                call_tree_id=call_tree_id,
                name=node.name,
                file_path=node.file_path,
                line_number=node.line_number,
                language=node.language,
                node_type=node.node_type,
                signature=node.signature,
                documentation=node.documentation,
                call_frequency=node.call_frequency,
                call_context=node.call_context,
                short_description=node.short_description,
                llm_explanation=node.llm_explanation,
                semantic_tags=node.semantic_tags,
                complexity_score=node.complexity_score,
                metadata={
                    "created_at": node.created_at.isoformat() if node.created_at else None,
                    "updated_at": node.updated_at.isoformat() if node.updated_at else None
                }
            )
            
            # Save the node
            node_model = await self.db_manager.create(node_model)
            self.stats["nodes_serialized"] += 1
            
            return node_model
            
        except Exception as e:
            self.stats["serialization_errors"] += 1
            error_response = handle_error(e, {
                "operation": "serialize_call_tree_node",
                "node_name": node.name,
                "file_path": node.file_path,
                "line_number": node.line_number
            })
            raise SerializationError(f"Failed to serialize call tree node: {error_response['message']}")
    
    async def deserialize_call_tree(self, tree_model: CallTreeModel) -> CallTree:
        """
        Deserialize a CallTreeModel to in-memory CallTree object.
        
        Args:
            tree_model: Database model instance
            
        Returns:
            CallTree object
        """
        try:
            # Create call tree
            call_tree = CallTree()
            call_tree.metadata = tree_model.metadata or {}
            call_tree.metadata.update({
                "created_at": tree_model.created_at,
                "analysis_type": tree_model.analysis_type,
                "language": tree_model.language,
                "total_nodes": tree_model.total_nodes,
                "max_depth": tree_model.max_depth
            })
            
            # Get all nodes for this tree
            node_models = await self.db_manager.get_by_filter(
                CallTreeNodeModel, 
                call_tree_id=tree_model.id
            )
            
            # Create mapping of node model IDs to node objects
            node_mapping = {}
            
            # First pass: create all nodes
            for node_model in node_models:
                node = await self.deserialize_call_tree_node(node_model)
                node_id = call_tree.add_node(node)
                node_mapping[node_model.id] = node
            
            # Second pass: rebuild relationships
            for node_model in node_models:
                node = node_mapping[node_model.id]
                
                # Add caller relationships
                for caller_id in (node_model.caller_node_ids or []):
                    if caller_id in node_mapping:
                        caller_node = node_mapping[caller_id]
                        node.add_caller(caller_node)
                
                # Add callee relationships  
                for callee_id in (node_model.callee_node_ids or []):
                    if callee_id in node_mapping:
                        callee_node = node_mapping[callee_id]
                        node.add_callee(callee_node)
            
            # Set root node if available
            if tree_model.root_node_id and tree_model.root_node_id in node_mapping:
                call_tree.root_node = node_mapping[tree_model.root_node_id]
            
            self.stats["trees_deserialized"] += 1
            log_debug(f"Deserialized call tree with {len(node_models)} nodes")
            
            return call_tree
            
        except Exception as e:
            self.stats["serialization_errors"] += 1
            error_response = handle_error(e, {
                "operation": "deserialize_call_tree",
                "tree_id": tree_model.id,
                "analysis_type": tree_model.analysis_type
            })
            raise SerializationError(f"Failed to deserialize call tree: {error_response['message']}")
    
    async def deserialize_call_tree_node(self, node_model: CallTreeNodeModel) -> CallTreeNode:
        """
        Deserialize a CallTreeNodeModel to in-memory CallTreeNode object.
        
        Args:
            node_model: Database model instance
            
        Returns:
            CallTreeNode object
        """
        try:
            node = CallTreeNode(
                name=node_model.name,
                file_path=node_model.file_path,
                line_number=node_model.line_number,
                language=node_model.language,
                node_type=node_model.node_type,
                signature=node_model.signature,
                documentation=node_model.documentation
            )
            
            # Set additional attributes
            node.call_frequency = node_model.call_frequency
            node.call_context = node_model.call_context
            node.short_description = node_model.short_description
            node.llm_explanation = node_model.llm_explanation
            node.semantic_tags = node_model.semantic_tags or []
            node.complexity_score = node_model.complexity_score
            
            # Set timestamps from metadata or model
            metadata = node_model.metadata or {}
            if "created_at" in metadata:
                node.created_at = datetime.fromisoformat(metadata["created_at"])
            else:
                node.created_at = node_model.created_at
                
            if "updated_at" in metadata:
                node.updated_at = datetime.fromisoformat(metadata["updated_at"])
            else:
                node.updated_at = node_model.updated_at
            
            self.stats["nodes_deserialized"] += 1
            
            return node
            
        except Exception as e:
            self.stats["serialization_errors"] += 1
            error_response = handle_error(e, {
                "operation": "deserialize_call_tree_node",
                "node_id": node_model.id,
                "node_name": node_model.name
            })
            raise SerializationError(f"Failed to deserialize call tree node: {error_response['message']}")
    
    def export_to_json(self, call_tree: CallTree, 
                      include_metadata: bool = True) -> str:
        """
        Export call tree to JSON format.
        
        Args:
            call_tree: CallTree to export
            include_metadata: Whether to include metadata
            
        Returns:
            JSON string representation
        """
        try:
            tree_dict = call_tree.to_dict()
            
            if not include_metadata:
                # Remove metadata fields
                tree_dict.pop("metadata", None)
                for node_dict in tree_dict.get("nodes", {}).values():
                    node_dict.pop("created_at", None)
                    node_dict.pop("updated_at", None)
            
            return json.dumps(tree_dict, indent=2, default=str)
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "export_to_json"})
            raise SerializationError(f"Failed to export to JSON: {error_response['message']}")
    
    def import_from_json(self, json_str: str) -> CallTree:
        """
        Import call tree from JSON format.
        
        Args:
            json_str: JSON string representation
            
        Returns:
            CallTree object
        """
        try:
            tree_dict = json.loads(json_str)
            
            # Create call tree
            call_tree = CallTree()
            call_tree.metadata = tree_dict.get("metadata", {})
            
            # Create nodes
            nodes_dict = tree_dict.get("nodes", {})
            node_mapping = {}
            
            # First pass: create all nodes
            for node_id, node_dict in nodes_dict.items():
                node = CallTreeNode(
                    name=node_dict["name"],
                    file_path=node_dict["file_path"],
                    line_number=node_dict["line_number"],
                    language=node_dict["language"],
                    node_type=node_dict.get("node_type", "function"),
                    signature=node_dict.get("signature"),
                    documentation=node_dict.get("documentation")
                )
                
                # Set additional attributes
                for attr in ["call_frequency", "call_context", "short_description", 
                           "llm_explanation", "semantic_tags", "complexity_score"]:
                    if attr in node_dict:
                        setattr(node, attr, node_dict[attr])
                
                # Set timestamps
                if "created_at" in node_dict:
                    node.created_at = datetime.fromisoformat(node_dict["created_at"])
                if "updated_at" in node_dict:
                    node.updated_at = datetime.fromisoformat(node_dict["updated_at"])
                
                call_tree.add_node(node)
                node_mapping[node_id] = node
            
            # Second pass: rebuild relationships
            for node_id, node_dict in nodes_dict.items():
                node = node_mapping[node_id]
                
                # Note: In JSON export, we don't store caller/callee relationships
                # as they are rebuilt during analysis
            
            # Set root node
            root_node_dict = tree_dict.get("root_node")
            if root_node_dict:
                root_node_id = f"{root_node_dict['file_path']}:{root_node_dict['line_number']}:{root_node_dict['name']}"
                if root_node_id in node_mapping:
                    call_tree.root_node = node_mapping[root_node_id]
            
            return call_tree
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "import_from_json"})
            raise SerializationError(f"Failed to import from JSON: {error_response['message']}")
    
    def export_to_binary(self, call_tree: CallTree, 
                        compress: bool = True) -> bytes:
        """
        Export call tree to binary format using pickle.
        
        Args:
            call_tree: CallTree to export
            compress: Whether to compress the data
            
        Returns:
            Binary data
        """
        try:
            data = pickle.dumps(call_tree, protocol=pickle.HIGHEST_PROTOCOL)
            
            if compress:
                data = gzip.compress(data)
            
            return data
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "export_to_binary"})
            raise SerializationError(f"Failed to export to binary: {error_response['message']}")
    
    def import_from_binary(self, binary_data: bytes, 
                          decompress: bool = True) -> CallTree:
        """
        Import call tree from binary format.
        
        Args:
            binary_data: Binary data
            decompress: Whether to decompress the data
            
        Returns:
            CallTree object
        """
        try:
            if decompress:
                binary_data = gzip.decompress(binary_data)
            
            call_tree = pickle.loads(binary_data)
            
            if not isinstance(call_tree, CallTree):
                raise SerializationError("Invalid binary data: not a CallTree object")
            
            return call_tree
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "import_from_binary"})
            raise SerializationError(f"Failed to import from binary: {error_response['message']}")
    
    def calculate_tree_hash(self, call_tree: CallTree) -> str:
        """
        Calculate a hash for the call tree for change detection.
        
        Args:
            call_tree: CallTree to hash
            
        Returns:
            SHA-256 hash string
        """
        try:
            # Create a normalized representation for hashing
            tree_data = {
                "nodes": {},
                "relationships": []
            }
            
            # Add nodes (sorted by ID for consistency)
            for node_id in sorted(call_tree.nodes.keys()):
                node = call_tree.nodes[node_id]
                tree_data["nodes"][node_id] = {
                    "name": node.name,
                    "file_path": node.file_path,
                    "line_number": node.line_number,
                    "language": node.language,
                    "node_type": node.node_type,
                    "signature": node.signature or ""
                }
            
            # Add relationships (sorted for consistency)
            relationships = []
            for node in call_tree.nodes.values():
                for callee in node.callee_nodes:
                    caller_id = f"{node.file_path}:{node.line_number}:{node.name}"
                    callee_id = f"{callee.file_path}:{callee.line_number}:{callee.name}"
                    relationships.append((caller_id, callee_id))
            
            tree_data["relationships"] = sorted(relationships)
            
            # Create hash
            tree_json = json.dumps(tree_data, sort_keys=True)
            return hashlib.sha256(tree_json.encode()).hexdigest()
            
        except Exception as e:
            error_response = handle_error(e, {"operation": "calculate_tree_hash"})
            raise SerializationError(f"Failed to calculate tree hash: {error_response['message']}")
    
    def get_serialization_stats(self) -> Dict[str, Any]:
        """Get serialization statistics."""
        return {
            "serialization_stats": self.stats.copy(),
            "database_stats": await self.db_manager.get_database_stats() if self.db_manager else {}
        }